#pragma once

// Single-level array
template <int BITS>
class TCMalloc_PageMap1
{
private:
	static const int LENGTH = 1 << BITS;
	void **array_;

public:
	typedef uintptr_t Number;

	// explicit TCMalloc_PageMap1(void* (*allocator)(size_t)) {
	explicit TCMalloc_PageMap1()
	{
		// array_ = reinterpret_cast<void**>((*allocator)(sizeof(void*) << BITS));
		size_t size = sizeof(void *) << BITS;
		size_t alignSize = SizeClass::_RoundUp(size, 1 << PAGE_SHIFT);
		array_ = (void **)SystemAlloc(alignSize >> PAGE_SHIFT);
		memset(array_, 0, sizeof(void *) << BITS);
	}

	// Return the current value for KEY.  Returns NULL if not yet set,
	// or if k is out of range.
	void *get(Number k) const
	{
		if ((k >> BITS) > 0)
		{
			return NULL;
		}
		return array_[k];
	}

	// REQUIRES "k" is in range "[0,2^BITS-1]".
	// REQUIRES "k" has been ensured before.
	//
	// Sets the value 'v' for key 'k'.
	void set(Number k, void *v)
	{
		array_[k] = v;
	}
};

// Two-level radix tree
template <int BITS>
class TCMalloc_PageMap2
{
private:
	// Put 32 entries in the root and (2^BITS)/32 entries in each leaf.
	static const int ROOT_BITS = 5;
	static const int ROOT_LENGTH = 1 << ROOT_BITS;

	static const int LEAF_BITS = BITS - ROOT_BITS;
	static const int LEAF_LENGTH = 1 << LEAF_BITS;

	// Leaf node
	struct Leaf
	{
		void *values[LEAF_LENGTH];
	};

	Leaf *root_[ROOT_LENGTH];	 // Pointers to 32 child nodes
	void *(*allocator_)(size_t); // Memory allocator

public:
	typedef uintptr_t Number;

	// explicit TCMalloc_PageMap2(void* (*allocator)(size_t)) {
	explicit TCMalloc_PageMap2()
	{
		// allocator_ = allocator;
		memset(root_, 0, sizeof(root_));

		PreallocateMoreMemory();
	}

	void *get(Number k) const
	{
		const Number i1 = k >> LEAF_BITS;
		const Number i2 = k & (LEAF_LENGTH - 1);
		if ((k >> BITS) > 0 || root_[i1] == NULL)
		{
			return NULL;
		}
		return root_[i1]->values[i2];
	}

	void set(Number k, void *v)
	{
		const Number i1 = k >> LEAF_BITS;
		const Number i2 = k & (LEAF_LENGTH - 1);
		assert(i1 < ROOT_LENGTH);
		root_[i1]->values[i2] = v;
	}

	bool Ensure(Number start, size_t n)
	{
		for (Number key = start; key <= start + n - 1;)
		{
			const Number i1 = key >> LEAF_BITS;

			// Check for overflow
			if (i1 >= ROOT_LENGTH)
				return false;

			// Make 2nd level node if necessary
			if (root_[i1] == NULL)
			{
				// Leaf* leaf = reinterpret_cast<Leaf*>((*allocator_)(sizeof(Leaf)));
				// if (leaf == NULL) return false;
				static ObjectPool<Leaf> leafPool;
				Leaf *leaf = (Leaf *)leafPool.New();
				if (leaf == nullptr)
					return false;

				memset(leaf, 0, sizeof(*leaf));
				root_[i1] = leaf;
			}

			// Advance key past whatever is covered by this leaf node
			key = ((key >> LEAF_BITS) + 1) << LEAF_BITS;
		}
		return true;
	}

	void PreallocateMoreMemory()
	{
		// Allocate enough to keep track of all possible pages
		Ensure(0, 1 << BITS);
	}
};

// three-level radix tree
template <int BITS>
class TCMalloc_PageMap3
{
private:
	typedef unsigned long long Number;
	// // 页大小配置 (8KB)
	// static const size_t kPageShift = 13;
	// static const size_t kPageSize = 1 << kPageShift;

	// 三级Radix Tree的位分配
	// 64位地址 = [Level1: 18位] [Level2: 18位] [Level3: 15位] [Offset: 13位]
	static const int kLevel1Bits = 18; // 高18位
	static const int kLevel2Bits = 18; // 中18位
	static const int kLevel3Bits = BITS - kLevel1Bits - kLevel2Bits;// 低15位,BITS已是位数减去PAGE_SHIFT结果

	static const size_t kLevel1Size = 1ULL << kLevel1Bits; // 256K 项
	static const size_t kLevel2Size = 1ULL << kLevel2Bits; // 256K 项
	static const size_t kLevel3Size = 1ULL << kLevel3Bits; // 32K 项
	
	//如果要进一步优化多线程并发情况,可以引入更加细粒度的锁:
	//可以在root中引入一把锁,在每一个Node2中引入一把锁
	//实际上,root中已经可以针对每一个Node2引入一把锁,但是直接分配锁就太多了,如果按需分配,那么针对每一个Node2锁的引入又需要锁进行保护,即需要一把大锁
	//此处的设计中,由于存在相关并发问题的PageMap使用,已经在外部的接口中加了一把大锁,故不再额外添加锁做进一步的线程并发优化

	// Level 3: 叶子节点，存储实际数据
	struct Leaf
	{
		void *values[kLevel3Size];

		Leaf()
		{
		}
	};

	// Level 2: 中间节点，指向Leaf
	struct Node2
	{
		Leaf *ptrs[kLevel2Size];

		Node2()
		{
			memset(ptrs, 0, sizeof(ptrs));
		}

		~Node2()
		{
			for (size_t i = 0; i < kLevel2Size; i++)
			{
				leafPool.Delete(ptrs[i]);
			}
		}
	};

	// // Level 1: 根节点，指向Node2
	struct Root
	{
		Node2 *ptrs[kLevel1Size];

		Root()
		{
			memset(ptrs, 0, sizeof(ptrs));
		}

		~Root()
		{
			for (size_t i = 0; i < kLevel1Size; i++)
			{
				node2Pool.Delete(ptrs[i]);
			}
		}
	};

	Root* root_; //直接对象,不使用指针,是无法做到的,因为栈帧空间没有那么大
	static inline ObjectPool<Node2> node2Pool;//用于node2的定长内存池
	static inline ObjectPool<Leaf> leafPool;//用于leaf的定长内存池
	static inline ObjectPool<Root> rootPool;//用于Root的定长内存池


	// 从页号中提取各级索引
	static inline size_t Level1Index(Number page_id)
	{
		return (page_id >> (kLevel2Bits + kLevel3Bits)) & ((1ULL << kLevel1Bits) - 1);
	}

	static inline size_t Level2Index(Number page_id)
	{
		return (page_id >> kLevel3Bits) & ((1ULL << kLevel2Bits) - 1);
	}

	static inline size_t Level3Index(Number page_id)
	{
		return page_id & ((1ULL << kLevel3Bits) - 1);
	}

public:
	TCMalloc_PageMap3()
	{
		// root_ = new Root();
		root_ = rootPool.New();
	}

	~TCMalloc_PageMap3()
	{
		// delete root_;
		rootPool.Delete(root_);
	}


	// 禁止拷贝和赋值
	TCMalloc_PageMap3(const TCMalloc_PageMap3 &) = delete;
	TCMalloc_PageMap3 &operator=(const TCMalloc_PageMap3 &) = delete;

	bool set(Number page_id, void* addr)
	{
		size_t idx1 = Level1Index(page_id);
		size_t idx2 = Level2Index(page_id);
		size_t idx3 = Level3Index(page_id);

		// 按需分配 Level 2 节点
		if (root_->ptrs[idx1] == nullptr)
		{
			root_->ptrs[idx1] = (Node2*)node2Pool.New();
			if (root_->ptrs[idx1] == nullptr)
			{
				return false; // 内存分配失败
			}
		}

		// 按需分配 Level 3 叶子节点
		Node2 *node2 = root_->ptrs[idx1];
		if (node2->ptrs[idx2] == nullptr)
		{
			node2->ptrs[idx2] = (Leaf*)leafPool.New();
			if (node2->ptrs[idx2] == nullptr)
			{
				return false; // 内存分配失败
			}
		}

		// 设置值
		Leaf *leaf = node2->ptrs[idx2];
		leaf->values[idx3] = addr;

		return true;
	}

	void* get(Number page_id)
	{
		size_t idx1 = Level1Index(page_id);
		size_t idx2 = Level2Index(page_id);
		size_t idx3 = Level3Index(page_id);

		// 逐级检查节点是否存在
		Node2 *node2 = root_->ptrs[idx1];
		if (node2 == nullptr)
			return nullptr;

		Leaf *leaf = node2->ptrs[idx2];
		if (leaf == nullptr)
			return nullptr;

		return leaf->values[idx3];
	}

	const void *get(Number page_id) const
	{
		return const_cast<TCMalloc_PageMap3 *>(this)->get(page_id);
	}

	void remove(Number page_id)
	{
		set(page_id,nullptr);
	}

	bool exists(Number page_id) const
	{
		return get(page_id) != nullptr;
	}

	bool exists(Number page_id)
	{
		return get(page_id) != nullptr;
	}

    struct Stats {
        size_t level2_nodes;  // Level 2 节点数
        size_t leaf_nodes;    // Leaf 节点数
        size_t memory_bytes;  // 总内存占用（字节）
    };
    
    Stats GetStats() const {
        Stats stats = {0, 0, 0};
        
        stats.memory_bytes += sizeof(Root);
        
        for (size_t i = 0; i < kLevel1Size; i++) {
            if (root_->ptrs[i] != nullptr) {
                stats.level2_nodes++;
                stats.memory_bytes += sizeof(Node2);
                
                Node2* node2 = root_->ptrs[i];
                for (size_t j = 0; j < kLevel2Size; j++) {
                    if (node2->ptrs[j] != nullptr) {
                        stats.leaf_nodes++;
                        stats.memory_bytes += sizeof(Leaf);
                    }
                }
            }
        }
        
        return stats;
    }
    
    /**
     * 打印统计信息
     */
    void PrintStats() const {
        Stats stats = GetStats();
        
        std::cout << "=== PageMap Statistics ===" << std::endl;
        std::cout << "Level 2 Nodes: " << stats.level2_nodes << std::endl;
        std::cout << "Leaf Nodes: " << stats.leaf_nodes << std::endl;
        std::cout << "Total Memory: " << stats.memory_bytes << " bytes ("
                  << stats.memory_bytes / 1024.0 / 1024.0 << " MB)" << std::endl;
        std::cout << "Theoretical Max (full tree): "
                  << (sizeof(Root) + kLevel1Size * sizeof(Node2) + 
                      kLevel1Size * kLevel2Size * sizeof(Leaf)) / 1024.0 / 1024.0 / 1024.0
                  << " GB" << std::endl;
    }
};
